from collections.abc import Iterable, Iterator
from typing import TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator

AnySample = TypeVar('AnySample')


class IterableDatasetIterator(DatasetIterator[AnySample]):
    def __init__(self, iterable: Iterable[AnySample]):
        self._iterator = iter(iterable)

    def __next__(self) -> AnySample:
        return next(self._iterator)


class IterableDataset(Dataset[AnySample]):
    def __init__(self, iterable: Iterable[AnySample]):
        if isinstance(iterable, Iterator):
            raise TypeError(f"Expected an reusable Iterable, got {iterable}")
        self._iterable = iterable

    @property
    def iterable(self) -> Iterable[AnySample]:
        return self._iterable

    def __len__(self) -> int:
        # noinspection PyTypeChecker
        return len(self._iterable)

    def __getitem__(self, index: int, /) -> AnySample:
        # noinspection PyUnresolvedReferences
        return self._iterable[index]

    def __iter__(self) -> IterableDatasetIterator[AnySample]:
        return IterableDatasetIterator(self._iterable)
